% vis_for_manuscript.m
%
% A subset of plots used in final manuscript figures. Most of these were
% generated in R, but TFRs and topoplots are created using FieldTrip.

saveFigs = false;

compname = getenv('computername');
if strcmp(compname,'DESKTOP-BHR0AU7')
    resultsdir = 'E:\ANL\Experiments\RESULTS\VAST\EEG\';
    figsavedir = 'E:\ANL\Experiments\RESULTS\VAST\FIGURES\TFRs\CNS\';
    addpath('C:\Users\jtjus\Documents\MATLAB\fieldtrip-20180512\fieldtrip-20180512');
    addpath('C:\Users\jtjus\Documents\MATLAB\DrosteEffect-BrewerMap-5b84f95');
    addpath(genpath('E:\MATLAB\CSDtoolbox'));
elseif strcmp(compname,'MSI')
    resultsdir = 'E:\ANL\RESULTS\VAST\EEG\';
    figsavedir = 'E:\ANL\RESULTS\VAST\FIGURES\TFRs\CNS\';
    addpath('D:\MATLAB\fieldtrip-20180809\fieldtrip-20180809');
    addpath('E:\MATLAB\DrosteEffect-BrewerMap-ca40391');
    addpath(genpath('D:\MATLAB\CSDtoolbox'));
else
    error('need to set up paths for this machine.')
end

%% Load and format TF data
load(strcat(resultsdir, "freq_band_timecourses.mat"))
% load(strcat(resultsdir, "freq_STATs.mat"))
temp = load(strcat(resultsdir, "avg_TFR.mat"));
avg_TFR = temp.avg_TFR;
avg_TFR_diff = temp.avg_TFR_diff;
clear temp
% Adjust timebases
NC = size(avg_TFR, 2);
NW = size(avg_TFR{2,1}, 2);
NC_diff = size(avg_TFR_diff, 2);
NW_diff = size(avg_TFR_diff{2,1}, 2);
trunc_enc_to = 2.05; % sec
INT_win_st = 0.5; % (sec) time after end ENC that INT win TFR starts
INT_task_st = 2; % (sec) time after end ENC of first INT task stimulus
min_len = Inf;
for c = 1:NC
    for w = 1:NW
        if avg_TFR{2,c}{2,w}.time(1) == 0 % not adjusted for 1 sec BL
            avg_TFR{2,c}{2,w}.time = avg_TFR{2,c}{2,w}.time - 1;
        end
    end
    % Find minimum length of the INT window across conditions
    if length(avg_TFR{2,c}{2,2}.time) < min_len
        min_len = length(avg_TFR{2,c}{2,2}.time);
    end
end
% Pass 2: truncate INT windows
for c = 1:NC
    avg_TFR{2,c}{2,2}.time = avg_TFR{2,c}{2,2}.time(1:min_len);
    avg_TFR{2,c}{2,2}.powspctrm = avg_TFR{2,c}{2,2}.powspctrm(:,:,1:min_len);
    % While here, convert to log scale (dB); already expressed re. baseline
    for win = 1:3 % encoding, retention, probe
        avg_TFR{2,c}{2,win}.powspctrm = 10 .* log10(avg_TFR{2,c}{2,win}.powspctrm);
    end
end

% Repeat for the difference TFRs
for c = 1:NC_diff
    for w = 1:NW_diff
        if avg_TFR_diff{2,c}{2,w}.time(1) == 0 % not adjusted for 1 sec BL
            avg_TFR_diff{2,c}{2,w}.time = avg_TFR_diff{2,c}{2,w}.time - 1;
        end
    end
    % Find minimum length of the INT window across conditions
    if length(avg_TFR_diff{2,c}{2,2}.time) < min_len
        min_len = length(avg_TFR_diff{2,c}{2,2}.time);
    end
end
% Pass 2: truncate INT windows
for c = 1:NC_diff
    avg_TFR_diff{2,c}{2,2}.time = avg_TFR_diff{2,c}{2,2}.time(1:min_len);
    avg_TFR_diff{2,c}{2,2}.powspctrm = avg_TFR_diff{2,c}{2,2}.powspctrm(:,:,1:min_len);
end

%% General plotting parameters
rbmap = colormap(flipud(brewermap(64,'RdBu'))); % blue, white, red
tfrcfg = [];
tfrcfg.comment = 'no';
tfrcfg.colormap = rbmap;
tfrcfg.colorbar = 'no';
tfrcfg.interactive = 'no';
tfrcfg.parameter = 'powspctrm';

topocfg = [];
topocfg.parameter = 'avg';
topocfg.comment = 'no';
topocfg.layout = 'biosemi64.lay';
topocfg.colormap = rbmap;
topocfg.colorbar = 'no';

INT_types = {'as', 'at', 'none'};

%% Alpha power topographies during WM encoding
for c = 1:NC % convert to log scale
    alpha_enc_mean{2,c}.avg = 10 .* log10(alpha_enc_mean{2,c}.avg);
end
topocfg.xlim = [alpha_enc_mean{2,1}.time(find(alpha_enc_mean{2,1}.time > 0, 1, 'first')),...
               alpha_enc_mean{2,1}.time(end)]; % already time restricted
% topocfg.zlim = [0.5 1.5]; % ratio scale
topocfg.zlim = [-1.8, 1.8]; % dB scale
topocfg.highlight = 'off';
% Average across WM domain (no differences in TFRs)
avg_byINT = [INT_types; cell(1, length(INT_types))];
for c = 1:length(INT_types)
    % aud WM
    avg_byINT{2,c} = alpha_enc_mean{2, contains(alpha_enc_mean(1,:), strcat('as_',INT_types{c}))};
    avg_byINT{2,c}.avg = squeeze(mean(cat(3, avg_byINT{2,c}.avg,...
        alpha_enc_mean{2, contains(alpha_enc_mean(1,:), strcat('at_',INT_types{c}))}.avg), 3));
    figure
    ft_topoplotER(topocfg, avg_byINT{2,c})
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'alpha_durENC_parietal_audWM_',INT_types{c},'INT.svg'))
        close all
    end
    
    % vis WM
    avg_byINT{2,c} = alpha_enc_mean{2, contains(alpha_enc_mean(1,:), strcat('vs_',INT_types{c}))};
    avg_byINT{2,c}.avg = squeeze(mean(cat(3, avg_byINT{2,c}.avg,...
        alpha_enc_mean{2, contains(alpha_enc_mean(1,:), strcat('vt_',INT_types{c}))}.avg), 3));
    figure
    ft_topoplotER(topocfg, avg_byINT{2,c})
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'alpha_durENC_parietal_visWM_',INT_types{c},'INT.svg'))
        close all
    end
end

%% TFRs through memory retention window (all conds, collapsed acx WM domain)
tfrcfg.channel = 'Pz';
% tfrcfg.channel = {'Pz','P1','CP3'}; % in response to CBPT results
tfrcfg.xlim = [avg_TFR{2,1}{2,2}.time(find(avg_TFR{2,1}{2,2}.time > 0, 1, 'first')),...
               avg_TFR{2,1}{2,2}.time(end)];
tfrcfg.ylim = [1, 30];
% tfrcfg.zlim = [0.5 1.5]; % ratio scale
tfrcfg.zlim = [-1.8, 1.8]; % dB scale
% Average across WM domain (no differences in TFRs)
avg_byINT = [INT_types; cell(1, length(INT_types))];
for c = 1:length(INT_types)
    % Aud WM
    avg_byINT{2,c} = avg_TFR{2, contains(avg_TFR(1,:), strcat('as_',INT_types{c}))}{2, 2};
    avg_byINT{2,c}.powspctrm = squeeze(mean(cat(4, avg_byINT{2,c}.powspctrm,...
        avg_TFR{2, contains(avg_TFR(1,:), strcat('at_',INT_types{c}))}{2, 2}.powspctrm), 4));
    figure
    ft_singleplotTFR(tfrcfg, avg_byINT{2,c})
    pbaspect([3 1 1])
    line([1.5, 1.5], [1 30], 'LineWidth', 8, 'LineStyle', '--', 'Color', [0.1 0.1 0.1]) 
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'RET_Parietal_audWM_',INT_types{c},'INT.svg'))
        close all
    end
    % Vis WM
    avg_byINT{2,c} = avg_TFR{2, contains(avg_TFR(1,:), strcat('vs_',INT_types{c}))}{2, 2};
    avg_byINT{2,c}.powspctrm = squeeze(mean(cat(4, avg_byINT{2,c}.powspctrm,...
        avg_TFR{2, contains(avg_TFR(1,:), strcat('vt_',INT_types{c}))}{2, 2}.powspctrm), 4));
    figure
    ft_singleplotTFR(tfrcfg, avg_byINT{2,c})
    pbaspect([3 1 1])
    line([1.5, 1.5], [1 30], 'LineWidth', 8, 'LineStyle', '--', 'Color', [0.1 0.1 0.1]) 
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'RET_Parietal_visWM_',INT_types{c},'INT.svg'))
        close all
    end
end

%% Alpha power topographies
% Note, during INT task (through response) for aud; post-INT for vis

% Aud -- variable already restricted to aud WM conditions
for i = 1:4 % convert to log scale
    alpha_durINT_domCollapsed_mean{2,i}.avg =...
        10 .* log10(alpha_durINT_domCollapsed_mean{2,i}.avg);
end
topocfg.xlim = [alpha_durINT_domCollapsed_mean{2,1}.time(1),...
               alpha_durINT_domCollapsed_mean{2,1}.time(end)]; % already time restricted
% topocfg.zlim = [0.5 1.5]; % ratio scale
topocfg.zlim = [-1.8, 1.8]; % dB scale
topocfg.highlight = 'off';
for c = 1:length(INT_types)
    figure
    ft_topoplotER(topocfg,...
        alpha_durINT_domCollapsed_mean{2, contains(alpha_durINT_domCollapsed_mean(1,:), INT_types{c})})
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'alpha_topo_durINT_audWM_',INT_types{c},'.svg'))
        close all
    end
end

% Vis -- variable already restricted to vis WM conditions
for i = 1:3 % convert to log scale
    alpha_postINT_domCollapsed_mean{2,i}.avg =...
        10 .* log10(alpha_postINT_domCollapsed_mean{2,i}.avg);
end
topocfg.xlim = [alpha_postINT_domCollapsed_mean{2,1}.time(1),...
               alpha_postINT_domCollapsed_mean{2,1}.time(end)]; % already time restricted
% topocfg.zlim = [0.5 1.5]; % ratio scale
topocfg.zlim = [-1.8, 1.8]; % dB scale
topocfg.highlight = 'off';
for c = 1:length(INT_types)
    figure
    ft_topoplotER(topocfg,...
        alpha_postINT_domCollapsed_mean{2, contains(alpha_postINT_domCollapsed_mean(1,:), INT_types{c})})
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'alpha_topo_postINT_visWM_',INT_types{c},'.svg'))
        close all
    end
end

% %% Statistical results of alpha power comparisons (topoplots)
% domCol_plotOrder = {'at','none';'as','none';'at','as'}; % plot subtraction in each row
% for i = 1:size(domCol_plotOrder, 1)
%     % Reset the stat topoplot cfg
%     % topocfg.zlim = [-0.35 0.35]; % ratio scale
%     topocfg.zlim = [-1.8 1.8]; % dB scale
%     topocfg.highlightcolor = [0 0 0];
%     topocfg.highlightsize = 12; % default = 6
%     topocfg.comment = 'no';
%     
%     % --- AUDITORY WM: DURING int task, avg across WM domain ---
%     % Subtract the corresponding data
%     aconds = alpha_durINT_domCollapsed_mean(1,:);
%     toPlot = alpha_durINT_domCollapsed_mean{2,...
%         contains(aconds, domCol_plotOrder{i,1})};
%     toPlot.avg = toPlot.avg - alpha_durINT_domCollapsed_mean{2,...
%         contains(aconds, domCol_plotOrder{i,2})}.avg;
%     % Select a statistical comparison
%     switch i
%         case 1
%             stat = STAT_alpha_durINT.aAT_aNone;
%         case 2
%             stat = STAT_alpha_durINT.aAS_aNone;
%         case 3
%             stat = STAT_alpha_durINT.aAT_aAS;
%     end
%     % Positive cluster?
%     if isfield(stat, 'posclusters') && ~isempty(stat.posclusters) && stat.posclusters(1).prob < 0.05
%         p_clus = true;
%         topocfg.highlight = 'on';
%         clab = stat.posclusterslabelmat == 1; % first cluster only
%         topocfg.highlightchannel = find(sum(clab, 2) > 0); % all chans that
%             % are part of the cluster at any time point
%         signifTime = stat.time(sum(clab, 1) > 0); % all T pt with signif
%         topocfg.xlim = [signifTime(1) signifTime(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title(strcat(num2str(topocfg.xlim(1)), " to ", num2str(topocfg.xlim(end)),...
%             " sec")) % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_durINT_alpha_pos_a',...
%                 upper(domCol_plotOrder{i,1}),'_a',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     else
%         p_clus = false;
%     end
%     % Negative cluster?
%     if isfield(stat, 'negclusters') && ~isempty(stat.negclusters) && stat.negclusters(1).prob < 0.05
%         n_clus = true;
%         topocfg.highlight = 'on';
%         clab = stat.negclusterslabelmat == 1; % first cluster only
%         topocfg.highlightchannel = find(sum(clab, 2) > 0); % all chans that
%             % are part of the cluster at any time point
%         signifTime = stat.time(sum(clab, 1) > 0); % all T pt with signif
%         topocfg.xlim = [signifTime(1) signifTime(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title(strcat(num2str(topocfg.xlim(1)), " to ", num2str(topocfg.xlim(end)),...
%             " sec")) % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_durINT_alpha_neg_a',...
%                 upper(domCol_plotOrder{i,1}),'_a',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     else
%         n_clus = false;
%     end
%     % If neither, make a No Cluster topoplot
%     if p_clus == false && n_clus == false
%         topocfg.highlight = 'off';
%         topocfg.xlim = [stat.time(1) stat.time(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title("No Significance") % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_durINT_alpha_empty_a',...
%                 upper(domCol_plotOrder{i,1}),'_a',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     end
%     
%     % --- VISUAL WM: JUST AFTER int task, avg across WM domain ---
%     % Subtract the corresponding data
%     vconds = alpha_postINT_domCollapsed_mean(1,:);
%     toPlot = alpha_postINT_domCollapsed_mean{2,...
%         contains(vconds, domCol_plotOrder{i,1})};
%     toPlot.avg = toPlot.avg - alpha_postINT_domCollapsed_mean{2,...
%         contains(vconds, domCol_plotOrder{i,2})}.avg;
%     % Select a statistical comparison
%     switch i
%         case 1
%             stat = STAT_alpha_postINT.vAT_vnone;
%         case 2
%             stat = STAT_alpha_postINT.vAS_vnone;
%         case 3
%             stat = STAT_alpha_postINT.vAS_vAT;
%     end
%     % Positive cluster?
%     if isfield(stat, 'posclusters') && ~isempty(stat.posclusters) && stat.posclusters(1).prob < 0.05
%         p_clus = true;
%         topocfg.highlight = 'on';
%         clab = stat.posclusterslabelmat == 1; % first cluster only
%         topocfg.highlightchannel = find(sum(clab, 2) > 0); % all chans that
%             % are part of the cluster at any time point
%         signifTime = stat.time(sum(clab, 1) > 0); % all T pt with signif
%         topocfg.xlim = [signifTime(1) signifTime(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title(strcat(num2str(topocfg.xlim(1)), " to ", num2str(topocfg.xlim(end)),...
%             " sec")) % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_postINT_alpha_pos_v',...
%                 upper(domCol_plotOrder{i,1}),'_v',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     else
%         p_clus = false;
%     end
%     % Negative cluster?
%     if isfield(stat, 'negclusters') && ~isempty(stat.negclusters) && stat.negclusters(1).prob < 0.05
%         n_clus = true;
%         topocfg.highlight = 'on';
%         clab = stat.negclusterslabelmat == 1; % first cluster only
%         topocfg.highlightchannel = find(sum(clab, 2) > 0); % all chans that
%             % are part of the cluster at any time point
%         signifTime = stat.time(sum(clab, 1) > 0); % all T pt with signif
%         topocfg.xlim = [signifTime(1) signifTime(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title(strcat(num2str(topocfg.xlim(1)), " to ", num2str(topocfg.xlim(end)),...
%             " sec")) % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_postINT_alpha_neg_v',...
%                 upper(domCol_plotOrder{i,1}),'_v',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     else
%         n_clus = false;
%     end
%     % If neither, make a No Cluster topoplot
%     if p_clus == false && n_clus == false
%         topocfg.highlight = 'off';
%         topocfg.xlim = [stat.time(1) stat.time(end)];
%         figure
%         ft_topoplotER(topocfg, toPlot)
%         title("No Significance") % a note for Illustrator editing
%         if saveFigs
%             fig = gcf;
%             saveas(fig, strcat(figsavedir,'STAT_postINT_alpha_empty_v',...
%                 upper(domCol_plotOrder{i,1}),'_v',upper(domCol_plotOrder{i,2}),'.svg'))
%             close all
%         end  
%     end
% end

%% --- Supplemental figures ---

%% Alpha activity during WM encoding, all WM conditions
% TFRs
tfrcfg.channel = 'Pz';
tfrcfg.ylim = [1, 30];
% tfrcfg.zlim = [0.5 1.5]; % ratio scale
tfrcfg.zlim = [-1.8, 1.8]; % dB scale
for c = 1:size(avg_TFR, 2)
    % Set time range based on cond
    % Vis TFRs are longer than aud because the max interval is longer (580
    % vs 340 ms). The longest actual encoding window should be 1.22 sec for
    % aud, 1.94 sec for vis. Baselines have been corrected to negative
    % values already. NOTE: Each ENC win has a 500 ms trigger bugger on
    % either side of stimulus onset/offset. The time windows below are
    % calculated based on shifting the ideal window times to account for
    % these no-action time windows
    if contains(avg_TFR{1,c}, 'as_') || contains(avg_TFR{1,c}, 'at_')
        tfrcfg.xlim = [0.5 2.22-0.5];
    else
        tfrcfg.xlim = [0.5 2.94-0.5];
    end
    figure
    ft_singleplotTFR(tfrcfg, avg_TFR{2,c}{2,1}) % last 1 ind is ENC win
    pbaspect([3 1 1])
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'ENC_',tfrcfg.channel,'_',avg_TFR{1,c},'.svg'))
        close all
    end
end

%% Theta power throughout the WM retention window
tfrcfg.channel = 'AFz';
tfrcfg.xlim = [0 avg_TFR{2,1}{2,2}.time(end)]; % plot the whole TFR 
tfrcfg.ylim = [1, 20];
% tfrcfg.zlim = [0.5 1.5]; % ratio scale
tfrcfg.zlim = [-1.8, 1.8]; % dB scale

% Topos 
topocfg.xlim = [1.5 3]; % Actually 2-3.5 sec after window correction
% topocfg.zlim = [0.5 1.5]; % ratio scale
topocfg.zlim = [-1.8, 1.8]; % dB scale
topocfg.highlight = 'off';

% Convert theta time courses to dB scale
% SEPERATE VARIABLE FOR DEBUGGING
x = theta_wholeINT_mean;
for i = 1:size(x, 2)
   x{2,i}.avg = 10 .* log10(x{2,i}.avg);
end

cond_names_theta = x(1,:);
for c = 1:length(cond_names_theta)
    % --- TFRs ---
    figure
    ft_singleplotTFR(tfrcfg,...
        avg_TFR{2, contains(avg_TFR(1,:), cond_names_theta{c})}{2, 2})
    line([1.5, 1.5], [1 20], 'LineWidth', 3, 'LineStyle', '--', 'Color', [0.1 0.1 0.1])
    pbaspect([3 1 1])
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'RETtheta_',tfrcfg.channel,'_',...
                           cond_names_theta{c},'.svg'))
        close all
    end
    
    % --- Power timecourse topos ---
    figure
    ft_topoplotER(topocfg, x{2, c})
    title('')
    set(gca, 'YTickLabel', [])
    set(gca, 'XTickLabel', [])
    if saveFigs
        fig = gcf;
        saveas(fig, strcat(figsavedir,'RETtheta_topo_',...
                           cond_names_theta{c},'.svg'))
        close all
    end
end


